# trainer/main.py - 修复版本
from fastapi import FastAPI, HTTPException, BackgroundTasks
from pydantic import BaseModel
from typing import List, Dict, Optional
import asyncio
import torch
import os
from pathlib import Path
from datetime import datetime
import json
import numpy as np

# 初始化FastAPI应用
app = FastAPI(title="Onpharos Training Service")

# 训练状态管理
class TrainingState:
    def __init__(self):
        self.is_running = False
        self.current_cycle = 0
        self.total_cycles = 0
        self.progress = 0.0  # 0.0-100.0
        self.logs = []
        self.last_updated = None
        # 添加训练指标
        self.performance_metrics = {
            'overall_performance': 0.0,
            'average_reward': 0.0,
            'survival_rate': 0.0,
            'training_loss': 0.0
        }

training_state = TrainingState()

# 数据模型定义
class TrainingConfig(BaseModel):
    cycles: int = 100
    episodes: int = 10
    model_name: str = "lightweight_ppo"
    learning_rate: float = 0.0003
    batch_size: int = 64

class TrainingStatusResponse(BaseModel):
    is_running: bool
    current_cycle: int
    total_cycles: int
    progress: float
    last_updated: Optional[str]
    logs: List[str]
    performance_metrics: Dict[str, float]  # 新增性能指标

class ExperienceData(BaseModel):
    episode_id: str
    agent_actions: Dict[str, str]
    rewards: Dict[str, float]
    game_state: Dict

# 训练数据存储路径
TRAINING_DATA_DIR = Path("/app/training_data")
MODELS_DIR = Path("/app/models")
TRAINING_DATA_DIR.mkdir(exist_ok=True)
MODELS_DIR.mkdir(exist_ok=True)

# 模拟模型训练（增强版本）
class Trainer:
    @staticmethod
    async def train(config: TrainingConfig):
        training_state.is_running = True
        training_state.total_cycles = config.cycles
        training_state.current_cycle = 0
        training_state.progress = 0.0
        
        try:
            # 模拟训练循环
            for cycle in range(config.cycles):
                if not training_state.is_running:
                    break
                    
                # 记录日志
                log_msg = f"Cycle {cycle+1}/{config.cycles} started"
                training_state.logs.append(f"[{datetime.now()}] {log_msg}")
                training_state.current_cycle = cycle + 1
                training_state.progress = (cycle + 1) / config.cycles * 100
                training_state.last_updated = datetime.now().isoformat()
                
                # 模拟训练指标（随机生成模拟数据）
                training_state.performance_metrics = {
                    'overall_performance': np.random.uniform(0.3, 0.9),
                    'average_reward': np.random.uniform(50, 200),
                    'survival_rate': np.random.uniform(0.6, 0.95),
                    'training_loss': np.random.uniform(0.1, 0.5)
                }
                
                # 模拟每轮训练耗时
                await asyncio.sleep(2)
                
                # 保存中间结果
                if (cycle + 1) % 10 == 0:
                    checkpoint_path = MODELS_DIR / f"{config.model_name}_cycle_{cycle+1}.pth"
                    # 实际项目中应保存真实模型权重
                    torch.save({"cycle": cycle+1, "state_dict": {}}, checkpoint_path)
                    training_state.logs.append(f"[{datetime.now()}] Checkpoint saved to {checkpoint_path}")
            
            final_msg = "Training completed successfully" if training_state.is_running else "Training stopped"
            training_state.logs.append(f"[{datetime.now()}] {final_msg}")
            
        except Exception as e:
            error_msg = f"Training failed: {str(e)}"
            training_state.logs.append(f"[{datetime.now()}] {error_msg}")
            raise HTTPException(status_code=500, detail=error_msg)
        finally:
            training_state.is_running = False
            training_state.last_updated = datetime.now().isoformat()

# API端点定义 - 修复端点路径
@app.get("/", response_model=Dict[str, str])
async def root():
    return {"message": "Training service is running", "status": "healthy"}

# 修复API端点路径，添加/api前缀
@app.post("/api/training/start", response_model=Dict[str, str])
async def start_training(config: TrainingConfig, background_tasks: BackgroundTasks):
    """开始训练 - 修复端点路径"""
    if training_state.is_running:
        raise HTTPException(status_code=400, detail="Training is already running")
    
    # 清空历史日志
    training_state.logs = []
    background_tasks.add_task(Trainer.train, config)
    return {"message": "Training started", "config": config.dict()}

@app.post("/api/training/stop", response_model=Dict[str, str])
async def stop_training():
    """停止训练 - 修复端点路径"""
    if not training_state.is_running:
        raise HTTPException(status_code=400, detail="No training is running")
    
    training_state.is_running = False
    return {"message": "Training stopped successfully"}

@app.get("/api/training/status", response_model=TrainingStatusResponse)
async def get_training_status():
    """获取训练状态 - 修复端点路径"""
    return {
        "is_running": training_state.is_running,
        "current_cycle": training_state.current_cycle,
        "total_cycles": training_state.total_cycles,
        "progress": training_state.progress,
        "last_updated": training_state.last_updated,
        "logs": training_state.logs[-50:],  # 返回最近50条日志
        "performance_metrics": training_state.performance_metrics  # 新增性能指标
    }

@app.get("/api/training/metrics", response_model=Dict[str, float])
async def get_training_metrics():
    """获取训练指标 - 新增端点"""
    return training_state.performance_metrics

@app.get("/api/training/export", response_model=Dict[str, str])
async def export_training_data():
    """导出训练数据 - 新增端点"""
    try:
        export_data = {
            "training_state": {
                "is_running": training_state.is_running,
                "current_cycle": training_state.current_cycle,
                "total_cycles": training_state.total_cycles,
                "progress": training_state.progress,
                "performance_metrics": training_state.performance_metrics
            },
            "logs": training_state.logs,
            "export_time": datetime.now().isoformat()
        }
        
        export_path = TRAINING_DATA_DIR / "training_export.json"
        with open(export_path, "w", encoding="utf-8") as f:
            json.dump(export_data, f, ensure_ascii=False, indent=2)
            
        return {"message": f"Training data exported to {export_path}", "file_path": str(export_path)}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to export training data: {str(e)}")

@app.post("/experience/save", response_model=Dict[str, str])
async def save_experience(data: ExperienceData):
    """接收游戏引擎发送的经验数据"""
    try:
        # 按episode_id分文件存储
        exp_path = TRAINING_DATA_DIR / f"experience_{data.episode_id}.json"
        with open(exp_path, "w", encoding="utf-8") as f:
            json.dump(data.dict(), f, ensure_ascii=False, indent=2)
        return {"message": f"Experience saved to {exp_path}"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to save experience: {str(e)}")

@app.get("/models/list", response_model=List[str])
async def list_models():
    """列出已保存的模型"""
    return [f.name for f in MODELS_DIR.glob("*.pth") if f.is_file()]

# 健康检查端点
@app.get("/health")
async def health_check():
    """健康检查"""
    return {
        "status": "healthy",
        "service": "training",
        "timestamp": datetime.now().isoformat(),
        "training_running": training_state.is_running
    }

# 系统信息端点
@app.get("/api/system/info")
async def get_system_info():
    """获取系统信息"""
    return {
        "service": "onpharos-trainer",
        "version": "2.0.0",
        "status": "running" if training_state.is_running else "idle",
        "current_cycle": training_state.current_cycle,
        "gpu_available": torch.cuda.is_available(),
        "models_count": len([f for f in MODELS_DIR.glob("*.pth")])
    }